"""
created: 2023-03-26 14:05
author:qinguoming
description: 
1.  这个程序是一个计算模板,用于对3D线性四节点单元组成的结构的静态分析,输出节点位移,输出单元应力,单元应变,节点支反力等结果.
2.  节点自由度为: ux,uy,uz      
3.  节点和网格的编号必须从1开始,且连续.
4.  所有set_name,surface_name必须是小写
5.  计算结果输出到mat文件

程序支持的边界条件和加载:
1. 位移边界条件(m):
    定义格式:
        a. (bc_name : str , node_label_list : list[int] , node_dof : int ,displace_value : float)
        b. (bc_name : str , node_set_name : str , node_dof : int ,displace_value : float)
        note: node_dof = 1 or 2 or 3 , 分别代表 ux , uy , uz 
        e.g: 
            ("DisplacementBc",[1,2,3,...],1,0.0)
2. 节点集中力加载(N):
    定义格式:
        a. ("ConcentratedLoad",load_name : str , node_label_list : list[int] , node_dof : int , force_value : float)
        b. ("ConcentratedLoad",load_name : str , node_set_name : str , node_dof : int , force_value : float)
        e.g: 
            ("ConcentratedLoad","load_2000N","nset_name",2,2000.0)];
            ("ConcentratedLoad","load_2000N",[1,2,3,4] ,2,2000.0)];
3. 单元面压力加载(Pa):
    定义格式:
        a. ("PressureLoad" , load_name : str ,surface_set_name : str , pressure_value : float)
        b. ("PressureLoad" , load_name : str ,array_of_element_label_and_faceId : ndarray(m,2), pressure_value : float)
        e.g: 
            ("PressureLoad" , "pressure_10Kpa" , "surf-1", 10000.0)
            surf_data=np.array([[1,1],[2,3],[15,2],[5,1]...])
            ("PressureLoad" , "pressure_5pa" ,surf_data, 5.0)
4. 体积力加载(N/m^3):
    定义格式:
        a. ("BodyLoad" , load_name : str , el_set_name : str , load_direction : str , body_force_value : float)
        b. ("BodyLoad" , load_name : str , element_label_list : list[int] , load_direction : str , body_force_value : float)
        note : load_direction = 'x' or 'y' or 'z'
        e.g: 
            ("BodyLoad" , "gravity_x_10k" , "elset-1", 'x' ,10000.0)
            ("BodyLoad" , "gravity_y_5" ,element_label_list, 'y' , 5.0)
5. 面力加载(pa):
    定义格式:
        a. ("SurfaceLoad" , load_name : str ,surface_set_name : str , direction_unit_vector = ndarray (3,) , value : float)
        b. ("SurfaceLoad" , load_name : str ,array_of_element_label_and_faceId : ndarray(m,2) , direction_unit_vector = ndarray (3,) , value : float)
        e.g: 
            ("SurfaceLoad" , "pressure_10Kpa" , "surf-1", [1.0,0,0] ,10000.0)
            surf_data=np.array([[1,1],[2,3],[15,2],[5,1]...])
            ("SurfaceLoad" , "pressure_5pa" ,surf_data, [1,0,0], 5.0)
"""
import numpy as np
from EF2D.tools import *
from EF2D.Elements import C3D4
import scipy.io as sio
import matplotlib.pyplot as plt
from typing import List,Dict
import sys
import os
from EF2D.tools import *
from EF2D import *
# 脚本执行: \EF2D> L:\FEACODE\DevEnv\PortablePython310\App\Python\python.exe -m Scripts.Script_for_3d_C3D4_StaticAnalysis
initNumpyPrintOptions()
# -------------定义变量------------------
job_name="job-c3d4-ToppPressure5mpa"
# 网格文件
mesh_file="L:\\EF2D\\tests\\data\\c3d4_beam.inp"
# mesh_file="L:\\EF2D\\tests\\data\\Job-1.inp"
# 材料属性(SI-mm单位制)
# rho=7860e-9
E=212000
nu=0.288
# 约束:支持位移边界条件
def_bc=[("DisplacementBc","fix_nset",1,0.0),
        ("DisplacementBc","fix_nset",2,0.0),
        ("DisplacementBc","fix_nset",3,0.0)]
# 定义载荷(载荷类型,载荷参数)
# def_load=[("ConcentratedLoad","n1000N_z","force_n1000N_z",3,-1000.0),]
# def_load=[("ConcentratedLoad","n1000N_y","Set_force1000N",2,-1000.0),]
# def_load=[("BodyLoad" ,"body_force-1","body_force_gravity_nz",'z',-500.0)]
# def_load=[ ("SurfaceLoad" , "surface_load-1000pa" , "surface_load_tensile_1000N", [-1.0,0,0] ,1000.0)]
def_load=[("PressureLoad" , "pressure_500pa" , "pressure_load_500mpa", 5.0)]

#-////////求解设置//////////////////////
bcMethod=2  # 边界条件处理方法,为1,2,3,4分别是Solver.py的方法
export_result=True  # 是否输出结果
export_path="L:/EF2D/tests/result/"

# -------------导入网格------------------
node_data, element_data, nset_data, elset_data ,surf_data= read_abaqus_inp(mesh_file)
print("节点数:",node_data.shape[0])
print("单元数:",element_data.shape[0])
print("节点集:",nset_data.keys())
print("单元集:",elset_data.keys())
print("surface集合:",surf_data.keys())
print("边界条件:",[b[0] for b in def_bc])
print("载荷有:",[l[1] for l in def_load])
print("边界条件处理方法:",bcMethod)

# -------------建立网格------------------
node_num=node_data.shape[0]
element_num=element_data.shape[0]
flobal_dof_num=node_num*3
# 定义节点对象
node_dict:Dict[int,Node]={}
for i in range(node_num):
    if node_data[i,0] not in node_dict.keys():
        node_dict[node_data[i,0]]=Node(label=node_data[i,0],x=node_data[i,1],y=node_data[i,2],z=node_data[i,3])
# 定义单元对象
element_list:List[C3D4]=[]
for i in range(element_num):
    element_list.append(C3D4(label=element_data[i,0],
                                        node1=node_dict[element_data[i,1]],
                                        node2=node_dict[element_data[i,2]],
                                        node3=node_dict[element_data[i,3]],
                                        node4=node_dict[element_data[i,4]],
                                        E=E,nu=nu))
# 定义载荷和边界
bc_list=[]
for b in def_bc:
    if b[0]=="DisplacementBc":
        if isinstance(b[1],str) and b[1].lower() in nset_data.keys():
            bc_list.append(DisplacementBc(name=b[0],nodeLabels=nset_data[b[1].lower()],dofth=b[2],value=b[3]))
        if (isinstance(b[1],list) or isinstance(b[1],np.ndarray)):
            bc_list.append(DisplacementBc(*b))
load_list=[]
for l in def_load:
    if l[0]=="ConcentratedLoad":
        if isinstance(l[2],str) and l[2].lower() in nset_data.keys():
            load_list.append(ConcentratedLoad(name=l[1],
                                            nodeLabels=nset_data[l[2].lower()],
                                            dofth=l[3],
                                            value=l[4]))
        else:
            load_list.append(ConcentratedLoad(*l[1:]))
    elif l[0]=="PressureLoad":
        if isinstance(l[2],str) and l[2].lower() in surf_data.keys():
            load_list.append(PressureLoad(name=l[1],
                                        ElementLabels=surf_data[l[2].lower()][:,0],
                                        faceIds=surf_data[l[2].lower()][:,1],
                                        value=l[3]))
        else:
            load_list.append(PressureLoad(name=l[1],
                                        ElementLabels=l[2][:,0],
                                        faceIds=l[2][:,1],
                                        value=l[3]))
    elif l[0] == "BodyLoad":
        if isinstance(l[2],str) and l[2].lower() in elset_data.keys():
            load_list.append(BodyLoad(name=l[1],
                                    ElementLabels=elset_data[l[2].lower()],
                                    Direction=l[3],
                                    value=l[4]))
        else:
            load_list.append(BodyLoad(*l[1:]))
    elif l[0] == "SurfaceLoad":
        if isinstance(l[2],str) and l[2].lower() in surf_data.keys():
            load_list.append(SurfaceLoad(name=l[1],
                                        ElementLabels=surf_data[l[2].lower()][:,0],
                                        faceIds=surf_data[l[2].lower()][:,1],
                                        directionVector=l[3],
                                        value=l[4]))
        else:
            load_list.append(SurfaceLoad(name=l[1],
                                        ElementLabels=l[2][:,0],
                                        faceIds=l[2][:,1],
                                        directionVector=l[3],
                                        value=l[4]))
                                        
# -------------组装刚度矩阵-------------------------
# Struc=FEA(meshblocks=block_list,bcs=bc_list,loads=load_list)
KG=np.zeros((flobal_dof_num,flobal_dof_num))
for e in element_list:
    e_id=e.label
    KG[np.ix_(e.ElemDofIndexs,e.ElemDofIndexs)]+=e.Ke
# 边界条件处理
bc_dict={}
for bc in bc_list:
    if isinstance(bc,DisplacementBc):
        for n_label in bc.NodeLabelList:
            dof=1+get_ithdof_index(n_label,3,bc.Dofth)
            if dof not in bc_dict.keys():
                bc_dict[dof]=bc.Value
            else:
                bc_dict[dof]+=bc.Value
# key: 从1开始的自由度编号, value: 加载节点力值
eqf=np.zeros(flobal_dof_num)
for load in load_list:
    if isinstance(load,ConcentratedLoad):
        for n_label in load.nodeLabelList:
            dof=1+get_ithdof_index(n_label,3,load.dofth)
            eqf[dof]+=load.LoadValue
    if isinstance(load,BodyLoad):
        for e in element_list:
            if e.label not in load.ElementLabels:
                continue
            e_body_force=e.ProcessELoad2EqNodeForceBodyLoad(load)
            eqf[e.ElemDofIndexs]+=e_body_force[:,0]
    if isinstance(load,PressureLoad):
        for e in element_list:
            if e.label not in load.ElementLabels:
                continue
            e_pressure_force=e.ProcessELoad2EqNodeForcePressureLoad(load)
            eqf[e.ElemDofIndexs]+=e_pressure_force[:,0]
    if isinstance(load,SurfaceLoad):
        for e in element_list:
            if e.label not in load.ElementLabels:
                continue
            e_surface_force=e.ProcessELoad2EqNodeForceSurfaceLoad(load)
            eqf[e.ElemDofIndexs]+=e_surface_force[:,0]
# 找出eqF中非0值的索引和非零值,定义全局节点力字典{key:自由度索引(0-base), value:节点力}
# P_ind=eqf.nonzero()
# load_dict=dict(zip((P_ind[0]+1).tolist(),Fg[P_ind].tolist()))
load_dict={ind+1:eqf[ind] for ind in eqf.nonzero()[0] }
# ic(load_dict)
# -------------运行分析------------------
if bcMethod==1:
    ug,fg=Method1(KG,bc_dict,load_dict)
if bcMethod==2:
    ug,fg=Method2(KG,bc_dict,load_dict)
if bcMethod==3:
    ug,fg=Method3(KG,bc_dict,load_dict) # type: ignore
if bcMethod==4:
    ug,fg=Method4(KG,bc_dict,load_dict) # type: ignore


# -------------输出结果--------------------

el_stress=np.zeros((element_num,6))
el_strain=np.zeros((element_num,6))
el_mises=np.zeros((element_num,1))
# 计算应力、应变、节点位移、节点支反力
for i,e in enumerate(element_list):
    el_stress[i,:]=e.calStress(ug[e.ElemDofIndexs,0]).T
    el_strain[i,:]=e.calStrain(ug[e.ElemDofIndexs,0]).T
    sx,sy,sz,sxy,syz,sxz=el_stress[i,:]
    el_mises[i,0]=np.sqrt(0.5*((sx-sy)**2+(sy-sz)**2+(sz-sx)**2+6*(sxy**2+syz**2+sxz**2)))
# 计算ux,uy,uz,umag
node_U=np.concatenate((ug[::3],ug[1::3],ug[2::3],np.sqrt(ug[::3]**2+ug[1::3]**2+ug[2::3]**2)),axis=1)
# 计算fx,fy,fz,fmag
node_F=np.concatenate((fg[::3],fg[1::3],fg[2::3],np.sqrt(fg[::3]**2+fg[1::3]**2+fg[2::3]**2)),axis=1)
ic(node_U.shape)
ic(node_F.shape)
ic(el_stress.shape)
ic(el_strain.shape)
ic(max(node_U[:,1]))
ic(min(node_U[:,1]))
# 输出结果到mat文件,可以和matlab对接
mat_data={
    "Node_data":node_data,
    "connecticity_data":element_data,
    "vtk_type":element_list[0].VTK_TYPE,
    "node_num":int(node_num),
    "element_num":int(element_num),
    "node_dof_num":3,
    "node_displacement":node_U,
    "node_force":node_F,
    "element_stress":el_stress,
    "element_strain":el_strain,
    "element_mises":el_mises,
    "nset":nset_data,
    "elset":elset_data,
    "surf":surf_data,
    }
sio.savemat(export_path+job_name+'.mat',mat_data,appendmat=False)
print("结果已输出到文件:",export_path+job_name+'.mat')

